"""
audio_predictor.py - 音频分类预测脚本优化版

功能：
1. 加载训练好的音频分类模型
2. 对音频文件进行预处理（梅尔频谱提取）
3. 执行预测并返回类别概率
4. 支持GPU加速（如果可用）
5. 详细的日志和错误处理

优化点：
- 添加类型注解
- 增强错误处理
- 支持批处理预测
- 添加配置选项
- 更完整的日志记录

使用方法：
python a07_predict.py models/trained/model.pth data/raw/human/sample_1743692044.wav --num_classes 2 --top 1

"""

import torch
import librosa
import numpy as np
import logging
from pathlib import Path
from typing import Tuple, Optional, Union, List, Dict

from utils.audio_cnn import AudioCNN

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger(__name__)


class AudioPredictor:
    """音频分类预测器

    参数:
        model_path (str): 训练好的模型权重路径
        num_classes (int): 类别数量
        sample_rate (int): 音频采样率 (默认44100)
        n_mels (int): 梅尔频带数 (默认128)
        device (str): 计算设备 ('cuda'或'cpu')
    """

    def __init__(
            self,
            model_path: str,
            num_classes: int = 5,
            sample_rate: int = 44100,
            n_mels: int = 128,
            device: Optional[str] = None
    ):
        # 设置计算设备
        self.device = self._get_device(device)
        logger.info(f"使用计算设备: {self.device}")

        # 加载模型
        self.model = self._load_model(model_path, num_classes)

        # 音频处理参数
        self.sample_rate = sample_rate
        self.n_mels = n_mels

        logger.info(f"预测器初始化完成 - 采样率: {sample_rate} | 梅尔频带数: {n_mels}")

    def _get_device(self, device: Optional[str]) -> torch.device:
        """自动选择最佳计算设备"""
        if device is None:
            return torch.device("cuda" if torch.cuda.is_available() else "cpu")
        return torch.device(device)

    def _load_model(self, model_path: str, num_classes: int) -> torch.nn.Module:
        """加载训练好的模型"""
        try:
            model = AudioCNN(num_classes=num_classes)

            # 加载模型权重 (支持多GPU训练保存的模型)
            state_dict = torch.load(model_path, map_location=self.device)
            if all(k.startswith('module.') for k in state_dict.keys()):
                state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}

            model.load_state_dict(state_dict)
            model.to(self.device)
            model.eval()

            logger.info(f"成功加载模型: {model_path}")
            return model
        except Exception as e:
            logger.error(f"加载模型失败: {str(e)}")
            raise

    def preprocess(
            self,
            audio_path: Union[str, Path],
            duration: Optional[float] = None,
            offset: float = 0.0
    ) -> torch.Tensor:
        """音频预处理: 加载音频并提取梅尔频谱特征

        参数:
            audio_path: 音频文件路径
            duration: 截取音频时长(秒) (None表示使用完整音频)
            offset: 开始截取的偏移量(秒)

        返回:
            形状为 [1, 1, n_mels, time] 的张量
        """
        try:
            # 加载音频文件
            y, sr = librosa.load(
                audio_path,
                sr=self.sample_rate,
                mono=True,
                duration=duration,
                offset=offset
            )

            # 提取梅尔频谱
            mel = librosa.feature.melspectrogram(
                y=y,
                sr=sr,
                n_mels=self.n_mels,
                fmax=sr // 2
            )

            # 转换为dB尺度并标准化
            mel_db = librosa.power_to_db(mel, ref=np.max)

            # 转换为PyTorch张量并添加批次和通道维度
            tensor = torch.FloatTensor(mel_db).unsqueeze(0).unsqueeze(0)

            logger.debug(f"预处理完成 - 输入形状: {tensor.shape}")
            return tensor.to(self.device)

        except Exception as e:
            logger.error(f"预处理音频 {audio_path} 失败: {str(e)}")
            raise

    def predict(
            self,
            audio_path: Union[str, Path],
            return_top: Optional[int] = None
    ) -> Union[np.ndarray, Dict[str, float]]:
        """执行音频分类预测

        参数:
            audio_path: 音频文件路径
            return_top: 返回前N个最高概率的类别 (None返回所有)

        返回:
            如果return_top为None: 返回所有类别概率数组
            如果return_top>0: 返回前N个类别及其概率的字典
        """
        try:
            # 预处理音频
            inputs = self.preprocess(audio_path)

            # 执行预测
            with torch.no_grad():
                outputs = self.model(inputs)
                probs = torch.softmax(outputs, dim=1)
                probs = probs.cpu().numpy()[0]  # 转换为numpy数组

            logger.debug(f"预测完成 - 原始概率: {probs}")

            # 根据return_top参数格式化输出
            if return_top is not None and return_top > 0:
                # 获取前N个最高概率的索引
                top_indices = np.argsort(probs)[-return_top:][::-1]
                return {str(i): float(probs[i]) for i in top_indices}
            return probs

        except Exception as e:
            logger.error(f"预测失败: {str(e)}")
            raise

    def predict_batch(
            self,
            audio_paths: List[Union[str, Path]],
            return_top: Optional[int] = None
    ) -> List[Union[np.ndarray, Dict[str, float]]]:
        """批量预测多个音频文件

        参数:
            audio_paths: 多个音频文件路径列表
            return_top: 返回前N个最高概率的类别 (None返回所有)

        返回:
            每个音频的预测结果列表
        """
        results = []
        for path in audio_paths:
            try:
                results.append(self.predict(path, return_top))
            except Exception as e:
                logger.error(f"处理文件 {path} 失败: {str(e)}")
                results.append(None)
        return results


if __name__ == "__main__":
    import argparse

    # 解析命令行参数
    parser = argparse.ArgumentParser(description='音频分类预测')
    parser.add_argument('model_path', type=str, help='模型权重路径')
    parser.add_argument('audio_path', type=str, help='要预测的音频文件路径')
    parser.add_argument('--num_classes', type=int, default=5, help='类别数量')
    parser.add_argument('--top', type=int, default=3, help='显示前N个最高概率的类别')
    parser.add_argument('--device', type=str, default=None,
                        help='计算设备 (cuda/cpu)，默认自动选择')
    args = parser.parse_args()

    try:
        ## 命令行
        # 初始化预测器
        predictor = AudioPredictor(
            model_path=args.model_path,
            num_classes=args.num_classes,
            device=args.device
        )
        # 执行预测
        result = predictor.predict(args.audio_path, return_top=args.top)

        ## 基本使用 非命令行
        # predictor = AudioPredictor("models/trained/model.pth")
        # predictor = AudioPredictor(
        #     model_path="models/trained/model.pth",
        #     num_classes=2,
        #     sample_rate=48000,
        #     n_mels=256,
        #     device="cpu"
        # )
        # result = predictor.predict("data/raw/human/sample_1743692044.wav")
        # results = predictor.predict_batch(["data/raw/human/sample_1743692044.wav", "data/raw/shovel/sample_1743671102.wav"]) # 批量预测

        # 打印结果
        print("\n预测结果:")
        if isinstance(result, dict):
            for cls, prob in result.items():
                print(f"类别 {cls}: {prob:.2%}")
        else:
            print("所有类别概率:")
            for i, prob in enumerate(result):
                print(f"类别 {i}: {prob:.2%}")

    except Exception as e:
        logger.error(f"运行预测失败: {str(e)}")